#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void translate_main(CodeGen *codegen); // 将 main 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

// TODO: 翻译你优化后的IR为RISC-V目标代码
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明外部函数 */
    codegen->append_inst(".extern getint", ASMInstruction::Atrribute);
    codegen->append_inst(".extern putint", ASMInstruction::Atrribute);

    /* 声明 main 函数 */ 
    codegen->append_inst(".align 2", ASMInstruction::Atrribute);
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main 函数开始 */
    codegen->append_inst("main", ASMInstruction::Label);

    /* main 函数的 Prologue (序言) */
    codegen->append_inst("addi", {"sp", "sp", "-64"});
    codegen->append_inst("sd", {"ra", "sp", "56"});
    codegen->append_inst("sd", {"s0", "sp", "48"});
    codegen->append_inst("addi", {"s0", "sp", "64"});

    /* main 函数的 label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);
    codegen->append_inst("# %op3 = call i32 @getint()", ASMInstruction::Comment);
    codegen->append_inst("call", {"getint"});
    codegen->append_inst("mv", {"s1", "a0"});  // 保存getint的结果到s1 (c的值)
    codegen->append_inst("# br label %label4", ASMInstruction::Comment);
    codegen->append_inst("li", {"t0", "1"});   // 初始化a的值为1
    codegen->append_inst("li", {"t1", "1"});   // 初始化b的值为1
    codegen->append_inst("j", {".main_label4"});

    /* label4: 主循环 */
    codegen->append_inst(".main_label4", ASMInstruction::Label);
    codegen->append_inst("# %op5 = phi i32 [ 1, %label_entry ], [ %op25, %label21 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op_b = phi i32 [ 1, %label_entry ], [ %op_b_new, %label21 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op6 = icmp slt i32 %op5, 999999999", ASMInstruction::Comment);
    codegen->append_inst("li", {"t2", "999999999"});
    codegen->append_inst("slt", {"t3", "t0", "t2"});
    codegen->append_inst("# br i1 %op6, label %label7, label %label12", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t4", "zero", "0"});
    codegen->append_inst("blt", {"t4", "t3", ".main_label7"});
    codegen->append_inst("j", {".main_label12"});

    /* label7: 检查除法条件 */
    codegen->append_inst(".main_label7", ASMInstruction::Label);
    codegen->append_inst("# %op10 = srem i32 %op5, %op3", ASMInstruction::Comment);
    codegen->append_inst("rem", {"t2", "t0", "s1"});
    codegen->append_inst("# %op11 = icmp eq i32 %op10, 0", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t3", "zero", "0"});
    codegen->append_inst("# br i1 %op11, label %label15, label %label18", ASMInstruction::Comment);
    codegen->append_inst("beq", {"t2", "t3", ".main_label15"});
    codegen->append_inst("j", {".main_label18"});

    /* label12: 循环结束 */
    codegen->append_inst(".main_label12", ASMInstruction::Label);
    codegen->append_inst("# call void @putint(i32 %op3)", ASMInstruction::Comment);
    codegen->append_inst("mv", {"a0", "s1"});
    codegen->append_inst("call", {"putint"});
    codegen->append_inst("# ret i32 %op5", ASMInstruction::Comment);
    codegen->append_inst("mv", {"a0", "t0"});
    codegen->append_inst("j", {".main_exit"});

    /* label15: a%c==0的情况 */
    codegen->append_inst(".main_label15", ASMInstruction::Label);
    codegen->append_inst("# %op17 = add i32 %op5, 3", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t0", "t0", "3"});  // a = a + 3
    codegen->append_inst("li", {"t1", "411"});        // b = 411
    codegen->append_inst("# br label %label21", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label21"});

    /* label18: a%c!=0的情况 */
    codegen->append_inst(".main_label18", ASMInstruction::Label);
    codegen->append_inst("# %op20 = add i32 %op5, 2", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t0", "t0", "2"});  // a = a + 2
    codegen->append_inst("li", {"t1", "101"});        // b = 101
    codegen->append_inst("# br label %label21", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label21"});

    /* label21: 计算a=a+b%3并返回循环开始 */
    codegen->append_inst(".main_label21", ASMInstruction::Label);
    codegen->append_inst("# %op_a_temp = phi i32 [ %op17, %label15 ], [ %op20, %label18 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op_b_new = phi i32 [ 411, %label15 ], [ 101, %label18 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op24 = srem i32 %op_b_new, 3", ASMInstruction::Comment);
    codegen->append_inst("li", {"t2", "3"});
    codegen->append_inst("rem", {"t3", "t1", "t2"});  // t3 = b % 3
    codegen->append_inst("# %op25 = add i32 %op_a_temp, %op24", ASMInstruction::Comment);
    codegen->append_inst("add", {"t0", "t0", "t3"});  // a = a + (b % 3)
    codegen->append_inst("# br label %label4", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label4"});

    /* main 函数的 Epilogue (收尾) */
    codegen->append_inst(".main_exit", ASMInstruction::Label);
    codegen->append_inst("ld", {"ra", "sp", "56"});
    codegen->append_inst("ld", {"s0", "sp", "48"});
    codegen->append_inst("addi", {"sp", "sp", "64"});
    codegen->append_inst("ret");
}
